import os
import numpy as np
from PIL import Image
import torch
from torchvision import models, transforms
from elasticsearch7 import Elasticsearch, exceptions

# ---------------------- 全局配置（尺寸在这里调整） ----------------------
TARGET_SIZE = (224, 224)  # 全局尺寸变量，修改此处即可调整所有尺寸
MODEL_IMAGE_DIMS = 2048    # ResNet50输出维度（与尺寸无关，模型固定）

# ---------------------- 图片预处理（全依赖全局尺寸） ----------------------
def load_image(image_path):
    image = Image.open(image_path)
    return image.convert('RGB') if image.mode == 'L' else image

def crop_to_content(image, tolerance=10):
    mask = image.getchannel('A') if image.mode == 'RGBA' else image.convert('L')
    mask_np = np.array(mask)
    non_bg = np.where(mask_np > tolerance)
    if not non_bg[0].size: return image
    return image.crop((non_bg[1].min(), non_bg[0].min(), non_bg[1].max()+1, non_bg[0].max()+1))

def resize_keep_ratio(image):
    """按全局尺寸的「最大边」比例缩放（无需传参，全依赖全局变量）"""
    target_max_edge = max(TARGET_SIZE)  # 取全局尺寸的最大边（如224×224的max是224）
    scale = target_max_edge / max(image.size)
    new_size = (int(image.size[0] * scale), int(image.size[1] * scale))
    return image.resize(new_size, Image.Resampling.LANCZOS)

def pad_to_target(image):
    """按全局尺寸居中填充（无需传参）"""
    pad_left = (TARGET_SIZE[0] - image.size[0]) // 2
    pad_top = (TARGET_SIZE[1] - image.size[1]) // 2
    padded_image = Image.new(
        image.mode, 
        TARGET_SIZE, 
        (0, 0, 0, 0) if image.mode == 'RGBA' else (255, 255, 255)
    )
    padded_image.paste(
        image, 
        (pad_left, pad_top), 
        mask=image if image.mode == 'RGBA' else None
    )
    return padded_image

def normalize_image(tensor):
    return transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(tensor)

# ---------------------- 模型与数据库初始化 ----------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_model = models.resnet50(pretrained=True).eval().to(device)
es = Elasticsearch(["http://localhost:9200"])
INDEX = "trademarks"

def init_index():
    if not es.indices.exists(INDEX):
        es.indices.create(index=INDEX, body={
            "mappings": {
                "properties": {
                    "image_vector": {"type": "dense_vector", "dims": MODEL_IMAGE_DIMS, "similarity": "cosine"},
                    "trademark_id": {"type": "keyword"},#此处建议使用regno也就是注册id
                    "image_path": {"type": "keyword"},
                    # ---------------------- 新增brand_info的结构定义 ----------------------
                    "brand_info": {
                        "type": "object",  # 嵌套对象类型
                        "properties": {
                            "name": {"type": "text", "analyzer": "ik_max_word"},  # 中文分词
                            "regno": {"type": "keyword"},  # 注册号（精确匹配）
                            "classid": {"type": "keyword"},  # 分类号（精确匹配）
                            "appdate": {"type": "date", "format": "yyyy-MM-dd"},  # 申请日期（日期类型）
                            "firsttrialno": {"type": "keyword"},  # 初审号
                            "firsttrialdate": {"type": "date", "format": "yyyy-MM-dd"},  # 初审日期
                            "announceno": {"type": "keyword"},  # 公告号
                            "announcedate": {"type": "date", "format": "yyyy-MM-dd"},  # 公告日期
                            "pic": {"type": "keyword"},  # 图片URL
                            "agent": {"type": "text", "analyzer": "ik_max_word"},  # 代理机构
                            "status": {"type": "text", "analyzer": "ik_max_word"},  # 状态
                            "registrant": {"type": "text", "analyzer": "ik_max_word"}  # 注册人
                        }
                    }
                    # -------------------------------------------------------------------
                }
            }
        })

# ---------------------- 特征提取与入库 ----------------------
def extract_image_feat(image_path):
    image = load_image(image_path)
    image = crop_to_content(image)
    image = resize_keep_ratio(image)    # 自动用全局尺寸的最大边缩放
    image = pad_to_target(image)        # 自动填充到全局尺寸
    tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
    tensor = normalize_image(tensor)
    with torch.no_grad():
        feat = image_model(tensor).squeeze().cpu().numpy()
    return feat / np.linalg.norm(feat)  # 归一化

def index_image(trademark_id, image_path):
    try:
        feat = extract_image_feat(image_path)
        es.index(index=INDEX, id=trademark_id, body={
            "trademark_id": trademark_id,
            "image_path": image_path,
            "image_vector": feat.tolist(),
            "brand_info":{
                "name": "手机工坊 SHOUJIDIY",
                "regno": "13145424",
                "classid": "14",
                "appdate": "2013-08-27",
                "firsttrialno": "1430",
                "firsttrialdate": "2014-11-06",
                "announceno": "1442",
                "announcedate": "2015-02-07",
                "pic": "http://api.jisuapi.com/trademark/upload/201807/31173651574910.jpg",
                "agent": "杭州龙华知识产权代理有限公司",
                "status": "注册申请完成",
                "registrant": "浙江富春江移动通信科技有限公司"
            }
            #此处添加商标其他信息
        })
        print(f"Indexed image: {trademark_id}")
    except Exception as e:
        print(f"Error: {trademark_id} - {str(e)}")

# ---------------------- 批量入库示例 ----------------------
if __name__ == "__main__":
    init_index()
    img_dir = "./trademark_imgs"
    for filename in os.listdir(img_dir):
        if filename.endswith(('.jpg', '.png')):
            trademark_id = os.path.splitext(filename)[0]#获取商标id，默认图片名称即为商标id。
            index_image(trademark_id, os.path.join(img_dir, filename))